package com.kingdee.eas.patch;

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import com.totoro.agent.JavaDynAgent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import sun.misc.Launcher;

import java.io.*;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.net.URL;
import java.net.URLDecoder;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Created by Niki on 2018/5/15 11:58
 */
public class AgentUtil {
    private static final Logger LOGGER = LoggerFactory.getLogger(AgentUtil.class);
    private static String classesPath;
    private static String agentParth;
    private static VirtualMachine virtualMachine;
    private static String pid;
    static {
        classesPath = AgentUtil.class.getClassLoader().getResource("").getPath();
        if (classesPath.contains("%20")) {
            classesPath = classesPath.replace("%20", " ");
        }
        LOGGER.info("类文件路径：{}", classesPath);
        agentParth = getJarPath();
        //当前进程
        String name = ManagementFactory.getRuntimeMXBean().getName();
        pid = name.split("@")[0];
        LOGGER.info("当前进程pid:{}", pid);
    }

    private static String getJarPath() {
        URL url = JavaDynAgent.class.getProtectionDomain().getCodeSource().getLocation();

        String filePath = null;
        String path = url.getPath();
        LOGGER.info("**********agent jar path is：{}", path);
        try {
            filePath = URLDecoder.decode(path, "utf-8");

        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }
        LOGGER.info("**********agent jar filePath is：{}", filePath);
        if (filePath.endsWith(".jar")) {
            File file = new File(filePath);

            return file.getAbsolutePath();
        } else {
            //无法获取到Agent的jar时通过写死来获得
            if (!Constant.useFixedUrl) {
                throw new RuntimeException("无法获取Agent jar包路径");
            }
            //该路径为服务器上agent的绝对路径
            filePath = Constant.agentPath;

            File file = new File(filePath);
            LOGGER.info("filePath 不符合规定！写死filePath is:{},返回其绝对路径为：{}", filePath, file.getAbsolutePath());
            return file.getAbsolutePath();
        }
    }

    private static void init() {
        try {
            virtualMachine = VirtualMachine.attach(pid);
            LOGGER.info("**** 代理路径为：{}", agentParth);
            virtualMachine.loadAgent(agentParth);
            virtualMachine.getAgentProperties();
        } catch (AttachNotSupportedException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (AgentLoadException e) {
            e.printStackTrace();
        } catch (AgentInitializationException e) {
            e.printStackTrace();
        }
        LOGGER.info("****虚拟机加载代理成功!");

        Instrumentation instrumentation = getInstrumentation();
        if (instrumentation == null) {
            throw new RuntimeException("initInstrumentation must not be null");
        }
    }

    public static void javaAgent(Map<String, String> classProperties) throws PatchException {
        if (classProperties == null || classProperties.isEmpty()) {
            return ;
        }
        init();
        List<ClassDefinition> classDefinitions = new ArrayList<>();
        for (String className : classProperties.keySet()) {
            try {
                Class<?> c = Class.forName(className);
                String classPath = classProperties.get(className);
                LOGGER.info("class redefined:{}", classPath);

                byte[] bytesFromFile = IOUtils.getBytes(classPath);
                ClassDefinition classDefinition = new ClassDefinition(c, bytesFromFile);
                classDefinitions.add(classDefinition);

                //覆盖原来的class文件
                coverOldClass(className, classPath);
            } catch (ClassNotFoundException e) {
                LOGGER.error("没有找到补丁类：{}",className,e);
                throw new PatchException("没找到补丁类：" + className);
            } finally {
                try {
                    destroy();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        try {
            getInstrumentation().redefineClasses(classDefinitions.toArray(new ClassDefinition[classDefinitions.size()]));
        } catch (ClassNotFoundException |UnmodifiableClassException e) {
            LOGGER.error("动态加载补丁异常！",e);
            throw new PatchException("动态加在补丁异常！");
        }
    }

    private static Instrumentation getInstrumentation() {
        try {
            Class clazz = Launcher.getLauncher().getClassLoader().loadClass(JavaDynAgent.class.getName());
            Field field = clazz.getDeclaredField("instrumentation");
            field.setAccessible(true);
            Instrumentation obj = (Instrumentation) field.get(clazz);
            return obj;
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } catch (NoSuchFieldException e) {
            e.printStackTrace();
        } catch (IllegalAccessException e) {
            e.printStackTrace();
        }
        return null;
    }

    private static void destroy() throws IOException {
        if (virtualMachine != null) {
            virtualMachine.detach();
        }
    }

    /**
     * 覆盖原来的class文件
     * @param className
     * @param newClassPath
     */
    private static void coverOldClass(String className, String newClassPath) {
        //覆盖之前的class文件
        String classPath_ = classesPath + className.replace(".", File.separator) + ".class";
        File classFile = new File(classPath_);
        if (!classFile.exists()) {
            return ;
        }
        FileOutputStream outputStream = null ;
        FileInputStream inputStream = null;
        try {
            outputStream = new FileOutputStream(classPath_);
            inputStream = new FileInputStream(newClassPath);
            IOUtils.copyLarge(inputStream, outputStream);
        } catch (IOException e) {
            e.printStackTrace();
        }finally {
            IOUtils.closeQuietly(outputStream);
            IOUtils.closeQuietly(inputStream);
        }
    }
}
